#include "gtest_mpi_listener.hpp"
#include "mpi/cipher_mpi_base.hpp"

#include <gtest/gtest.h>
#include <helib/helib.h>

int root = 0;

// Plaintext prime modulus
unsigned long p = 97;
// Cyclotomic polynomial - defines phi(m)
unsigned long m = 32;
// Hensel lifting (default = 1)
unsigned long r = 1;
// Number of bits of the modulus chain
unsigned long bits = 500;
// Number of columns of Key-Switching matrix (default = 2 or 3)
unsigned long c = 2;

helib::Context context =
    helib::ContextBuilder<helib::BGV>().m(m).p(p).r(r).bits(bits).c(c).build();
helib::SecKey *secret_key = nullptr;
helib::PubKey *public_key = nullptr;

TEST(CIPHER_MPI, IREDUCE_ADD) {
  // Get MPI metadata
  int rank, np;
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &np);

  ASSERT_GE(np, 4);

  // Get the EncryptedArray of the context
  const helib::EncryptedArray &ea = context.getEA();

  // Get the number of slot (phi(m))
  long nslots = ea.size();

  // Create a vector of long with nslots elements
  helib::Ptxt<helib::BGV> ptxt(context);
  for (int i = 0; i < ptxt.size(); ++i) {
    ptxt[i] = rank;
  }

  // Create a ciphertext object
  helib::Ctxt ctxt(*public_key);
  // Encrypt the plaintext using the public_key
  public_key->Encrypt(ctxt, ptxt);

  // Reduce ctxt
  helib::Ctxt *ctxt_res = nullptr;
  if (rank == root) {
    ctxt_res = new helib::Ctxt(*public_key);
  }
  ctxt_ireduce_add(ctxt, ctxt_res, root, MPI_COMM_WORLD);

  if (rank == root) {
    // Create a plaintext for decryption
    helib::Ptxt<helib::BGV> plaintext_result(context);
    // Decrypt the modified ciphertext
    secret_key->Decrypt(plaintext_result, *ctxt_res);

    auto expect = ((0 + np - 1) * np / 2) % p;
    for (auto i = 0; i < plaintext_result.size(); i++) {
      auto actual =
          NTL::conv<long>(NTL::coeff(plaintext_result[i].getData(), 0));
      ASSERT_EQ(expect, actual);
    }
  }

  if (rank == root) {
    delete ctxt_res;
  }
}

TEST(CIPHER_MPI, REDUCE_ADD) {
  // Get MPI metadata
  int rank, np;
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &np);

  ASSERT_GE(np, 4);

  // Get the EncryptedArray of the context
  const helib::EncryptedArray &ea = context.getEA();

  // Get the number of slot (phi(m))
  long nslots = ea.size();

  // Create a vector of long with nslots elements
  helib::Ptxt<helib::BGV> ptxt(context);
  for (int i = 0; i < ptxt.size(); ++i) {
    ptxt[i] = rank;
  }

  // Create a ciphertext object
  helib::Ctxt ctxt(*public_key);
  // Encrypt the plaintext using the public_key
  public_key->Encrypt(ctxt, ptxt);

  // Reduce ctxt
  helib::Ctxt *ctxt_res = nullptr;
  if (rank == root) {
    ctxt_res = new helib::Ctxt(*public_key);
  }
  ctxt_reduce_add(ctxt, ctxt_res, root, MPI_COMM_WORLD);

  if (rank == root) {
    // Create a plaintext for decryption
    helib::Ptxt<helib::BGV> plaintext_result(context);
    // Decrypt the modified ciphertext
    secret_key->Decrypt(plaintext_result, *ctxt_res);

    auto expect = ((0 + np - 1) * np / 2) % p;
    for (auto i = 0; i < plaintext_result.size(); i++) {
      auto actual =
          NTL::conv<long>(NTL::coeff(plaintext_result[i].getData(), 0));
      ASSERT_EQ(expect, actual);
    }
  }

  if (rank == root) {
    delete ctxt_res;
  }
}

TEST(CIPHER_MPI, ALLREDUCE_ADD) {
  // Get MPI metadata
  int rank, np;
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &np);

  ASSERT_GE(np, 4);

  // Get the EncryptedArray of the context
  const helib::EncryptedArray &ea = context.getEA();

  // Get the number of slot (phi(m))
  long nslots = ea.size();

  // Create a vector of long with nslots elements
  helib::Ptxt<helib::BGV> ptxt(context);
  for (int i = 0; i < ptxt.size(); ++i) {
    ptxt[i] = rank;
  }

  // Create a ciphertext object
  helib::Ctxt ctxt(*public_key);
  // Encrypt the plaintext using the public_key
  public_key->Encrypt(ctxt, ptxt);

  // Reduce ctxt
  helib::Ctxt *ctxt_res = nullptr;
  ctxt_res = new helib::Ctxt(*public_key);
  ctxt_allreduce_add(ctxt, ctxt_res, MPI_COMM_WORLD);

  if (rank == root) {
    // Create a plaintext for decryption
    helib::Ptxt<helib::BGV> plaintext_result(context);
    // Decrypt the modified ciphertext
    secret_key->Decrypt(plaintext_result, *ctxt_res);

    auto expect = ((0 + np - 1) * np / 2) % p;
    for (auto i = 0; i < plaintext_result.size(); i++) {
      auto actual =
          NTL::conv<long>(NTL::coeff(plaintext_result[i].getData(), 0));
      ASSERT_EQ(expect, actual);
    }
  }

  delete ctxt_res;
}

int main(int argc, char **argv) {
  ::testing::InitGoogleTest(&argc, argv);
  int ret = MPI_Init(&argc, &argv);
  EXPECT_EQ(ret, 0);

  // Add object that will finalize MPI on exit; Google Test owns this pointer
  ::testing::AddGlobalTestEnvironment(new GTestMPIListener::MPIEnvironment);

  // Get the event listener list.
  ::testing::TestEventListeners &listeners =
      ::testing::UnitTest::GetInstance()->listeners();

  // Remove default listener: the default printer and the default XML printer
  delete listeners.Release(listeners.default_result_printer());
  delete listeners.Release(listeners.default_xml_generator());

  // Adds MPI listener; Google Test owns this pointer
  listeners.Append(new GTestMPIListener::MPIMinimalistPrinter);

  // Get MPI metadata
  int rank;
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);

  if (rank == root) {
    context.printout();

    // Generate the secret key
    secret_key = new helib::SecKey(context);
    secret_key->GenSecKey();

    // Compute key-switching matrices that we need
    helib::addSome1DMatrices(*secret_key);

    public_key = secret_key;
  }
  // Broadcast public key
  bcast_pubkey(context, public_key, root, MPI_COMM_WORLD);

  // ! Only for debug
  // ! Broadcast Seckey
  // bcast_seckey(context, secret_key, root, MPI_COMM_WORLD);

  // Run Test
  ret = RUN_ALL_TESTS();
  EXPECT_EQ(ret, 0);

  if (rank != root) {
    delete public_key;
  }
  // delete secret_key;

  return 0;
}